/*
 * Copyright (c) 2015-2017, ARM Limited and Contributors. All rights reserved.
 *
 * SPDX-License-Identifier: BSD-3-Clause
 */

#include <getopt.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <openssl/conf.h>
#include <openssl/evp.h>
#include <openssl/pem.h>

#if USE_TBBR_DEFS
#include <tbbr_oid.h>
#else
#include <platform_oid.h>
#endif

#include "cert.h"
#include "cmd_opt.h"
#include "debug.h"
#include "key.h"
#include "sha.h"

#define MAX_FILENAME_LEN        1024

/*
 * Create a new key container
 */
int key_new(key_t *key)
{
    /* Create key pair container */
    key->key = EVP_PKEY_new();
    if (key->key == NULL) {
        return 0;
    }

    return 1;
}

static int key_create_rsa(key_t *key)
{
    RSA *rsa;

    rsa = RSA_generate_key(RSA_KEY_BITS, RSA_F4, NULL, NULL);
    if (rsa == NULL) {
        printf("Cannot create RSA key\n");
        goto err;
    }
    if (!EVP_PKEY_assign_RSA(key->key, rsa)) {
        printf("Cannot assign RSA key\n");
        goto err;
    }

    return 1;
err:
    RSA_free(rsa);
    return 0;
}

#ifndef OPENSSL_NO_EC
static int key_create_ecdsa(key_t *key)
{
    EC_KEY *ec;

    ec = EC_KEY_new_by_curve_name(NID_X9_62_prime256v1);
    if (ec == NULL) {
        printf("Cannot create EC key\n");
        goto err;
    }
    if (!EC_KEY_generate_key(ec)) {
        printf("Cannot generate EC key\n");
        goto err;
    }
    EC_KEY_set_flags(ec, EC_PKEY_NO_PARAMETERS);
    EC_KEY_set_asn1_flag(ec, OPENSSL_EC_NAMED_CURVE);
    if (!EVP_PKEY_assign_EC_KEY(key->key, ec)) {
        printf("Cannot assign EC key\n");
        goto err;
    }

    return 1;
err:
    EC_KEY_free(ec);
    return 0;
}
#endif /* OPENSSL_NO_EC */

typedef int (*key_create_fn_t)(key_t *key);
static const key_create_fn_t key_create_fn[KEY_ALG_MAX_NUM] = {
    key_create_rsa,
#ifndef OPENSSL_NO_EC
    key_create_ecdsa,
#endif /* OPENSSL_NO_EC */
};

int key_create(key_t *key, int type)
{
    if (type >= KEY_ALG_MAX_NUM) {
        printf("Invalid key type\n");
        return 0;
    }

    if (key_create_fn[type]) {
        return key_create_fn[type](key);
    }

    return 0;
}

int key_load(key_t *key, unsigned int *err_code)
{
    FILE *fp;
    EVP_PKEY *k;

    if (key->fn) {
        /* Load key from file */
        fp = fopen(key->fn, "r");
        if (fp) {
            k = PEM_read_PrivateKey(fp, &key->key, NULL, NULL);
            fclose(fp);
            if (k) {
                *err_code = KEY_ERR_NONE;
                return 1;
            } else {
                ERROR("Cannot load key from %s\n", key->fn);
                *err_code = KEY_ERR_LOAD;
            }
        } else {
            WARN("Cannot open file %s\n", key->fn);
            *err_code = KEY_ERR_OPEN;
        }
    } else {
        WARN("Key filename not specified\n");
        *err_code = KEY_ERR_FILENAME;
    }

    return 0;
}

int key_store(key_t *key)
{
    FILE *fp;

    if (key->fn) {
        fp = fopen(key->fn, "w");
        if (fp) {
            PEM_write_PrivateKey(fp, key->key,
                    NULL, NULL, 0, NULL, NULL);
            fclose(fp);
            return 1;
        } else {
            ERROR("Cannot create file %s\n", key->fn);
        }
    } else {
        ERROR("Key filename not specified\n");
    }

    return 0;
}

int key_init(void)
{
    cmd_opt_t cmd_opt;
    key_t *key;
    unsigned int i;

    for (i = 0; i < num_keys; i++) {
        key = &keys[i];
        if (key->opt != NULL) {
            cmd_opt.long_opt.name = key->opt;
            cmd_opt.long_opt.has_arg = required_argument;
            cmd_opt.long_opt.flag = NULL;
            cmd_opt.long_opt.val = CMD_OPT_KEY;
            cmd_opt.help_msg = key->help_msg;
            cmd_opt_add(&cmd_opt);
        }
    }

    return 0;
}

key_t *key_get_by_opt(const char *opt)
{
    key_t *key;
    unsigned int i;

    /* Sequential search. This is not a performance concern since the number
     * of keys is bounded and the code runs on a host machine */
    for (i = 0; i < num_keys; i++) {
        key = &keys[i];
        if (0 == strcmp(key->opt, opt)) {
            return key;
        }
    }

    return NULL;
}
